module aes_256_top(
  input           clk,        // 50MHz时钟（满足50MB/s吞吐量：50MHz×1字节=50MB/s）
  input           rst_n,      // 低电平复位
  input           encrypt_en, // 加密使能
  input [127:0]   plaintext,  // 明文（128位，AES分组长度）
  input [255:0]   key,        // 密钥（256位）
  output reg      done,       // 加密完成标志
  output reg [127:0] ciphertext // 密文（128位）
);

// 密钥扩展模块：256位密钥→15轮密钥（AES-256需14轮加密+初始轮）
wire [127:0] round_keys [0:14];
aes_key_expand u_key_expand(
  .clk(clk),
  .rst_n(rst_n),
  .key(key),
  .round_keys(round_keys)
);

// 流水线寄存器（每轮加密结果暂存）
reg [127:0] round_data [0:14];
reg [3:0]    round_cnt; // 轮数计数

// 初始轮：明文与初始轮密钥异或
always @(posedge clk or negedge rst_n) begin
  if (!rst_n) begin
    round_data[0] <= 128'd0;
  end else if (encrypt_en) begin
    round_data[0] <= plaintext ^ round_keys[0];
  end
end

// 第1~13轮加密（每轮含SubBytes、ShiftRows、MixColumns、AddRoundKey）
genvar i;
generate
  for (i=1; i<=13; i=i+1) begin : round_loop
    aes_round u_round(
      .clk(clk),
      .rst_n(rst_n),
      .data_in(round_data[i-1]),
      .round_key(round_keys[i]),
      .data_out(round_data[i])
    );
  end
endgenerate

// 第14轮加密（无MixColumns）
always @(posedge clk or negedge rst_n) begin
  if (!rst_n) begin
    round_data[14] <= 128'd0;
  end else begin
    // SubBytes + ShiftRows + AddRoundKey
    wire [127:0] sub_shift_data;
    aes_subbytes_shiftrows u_sub_shift(
      .data_in(round_data[13]),
      .data_out(sub_shift_data)
    );
    round_data[14] <= sub_shift_data ^ round_keys[14];
  end
end

// 加密完成标志与密文输出
always @(posedge clk or negedge rst_n) begin
  if (!rst_n) begin
    done <= 1'b0;
    ciphertext <= 128'd0;
    round_cnt <= 4'd0;
  end else if (encrypt_en) begin
    if (round_cnt == 4'd14) begin // 14轮完成
      done <= 1'b1;
      ciphertext <= round_data[14];
      round_cnt <= 4'd0;
    end else begin
      done <= 1'b0;
      round_cnt <= round_cnt + 4'd1;
    end
  end else begin
    done <= 1'b0;
    round_cnt <= 4'd0;
  end
end

endmodule